package com.zixi.ai.chat.service.impl;

import com.zixi.ai.ai.advice.ReasoningContentAdvisor;
import com.zixi.ai.ai.core.IQSSearchEngine;
import com.zixi.ai.ai.core.entity.DataClean;
import com.zixi.ai.ai.service.IAiModelService;
import com.zixi.ai.chat.dto.ChatResponseDto;
import com.zixi.ai.chat.service.IWebSearchService;
import com.zixi.ai.llm.entity.AiLlmEntity;
import com.zixi.ai.llm.service.IAiLlmService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.expansion.QueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import javax.annotation.PostConstruct;
import java.util.Optional;

/**
 * @author wfg
 */
@Slf4j
@Service("webSearchServiceImpl")
public class WebSearchServiceImpl implements IWebSearchService {
    @Autowired
    private IAiModelService aiModelService;
    @Autowired
    private IAiLlmService aiLlmService;

    @Autowired
    private SimpleLoggerAdvisor simpleLoggerAdvisor;
    @Autowired
    @Qualifier("queryArgumentPromptTemplate")
    private PromptTemplate queryArgumentPromptTemplate;

    @Autowired
    @Qualifier("transformerPromptTemplate")
    private PromptTemplate transformerPromptTemplate;

    @Autowired
    private IQSSearchEngine searchEngine;
    @Autowired
    private DataClean dataCleaner;
    private final ReasoningContentAdvisor reasoningContentAdvisor = new ReasoningContentAdvisor(1);

    private WebSearchRetriever webSearchRetriever;
    @PostConstruct
    public void init() {
        webSearchRetriever = WebSearchRetriever.builder()
                .searchEngine(searchEngine)
                .dataCleaner(dataCleaner)
                .maxResults(2)
                .build();
    }
    @Override
    public Flux<ChatResponseDto> chat(String prompt, Long aiLlmId) {
        Optional<AiLlmEntity> aiLlm = aiLlmService.findById(aiLlmId);
        ChatModel chatModel;
        if (aiLlm.isEmpty()) {
            log.error("chat model is null");
            return Flux.empty();
        }
        chatModel = aiModelService.buildChatModel(aiLlm.get().getOptions());
        ChatClient chatClient = ChatClient.builder(chatModel).build();
        return chatClient.prompt()
                .advisors(
                        createRetrievalAugmentationAdvisor(chatClient),
                        reasoningContentAdvisor,
                        simpleLoggerAdvisor
                ).user(prompt)
                .stream().chatResponse().map(chatResponse
                        -> new ChatResponseDto(chatResponse.getResult().getOutput().getText())
                );
    }
    private RetrievalAugmentationAdvisor createRetrievalAugmentationAdvisor(ChatClient chatClient) {

        QueryTransformer queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(chatClient.mutate())
                .promptTemplate(transformerPromptTemplate)
                .targetSearchSystem("Web Search")
                .build();
        QueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClient.mutate())
                .numberOfQueries(2)
                .build();
        return RetrievalAugmentationAdvisor.builder()
                .documentRetriever(webSearchRetriever)
                .queryTransformers(queryTransformer)
                .queryAugmenter(
                        new CustomContextQueryAugmenter(
                                queryArgumentPromptTemplate,
                                null,
                                true)
                ).queryExpander(queryExpander)
                .documentJoiner(new ConcatenationDocumentJoiner())
                .build();
    }
}
